#![no_std]
#![doc = "PMU (Power Management Unit) driver for Rockchip RK3588 SoC"]
#![allow(static_mut_refs)]

use core::{
    ptr::{read_volatile, write_volatile},
    sync::atomic::{Ordering, compiler_fence},
};

pub mod regs;
pub use regs::{PowerDomain, PowerState, RebootLevel};

/// Timeout loop count (approximately 10ms)
const TIMEOUT_LOOPS: u32 = 10000;

/// RK3588 PMU base address
pub const PMU_BASE_ADDR: usize = 0xFD8D8000;

/// RK3588 CRU base address
pub const CRU_BASE_ADDR: usize = 0xFD7C0000;

/// RK3588 PMU driver
pub struct RK3588Pmu {
    base_addr: usize,
}

/// PMU operation error types
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PmuError {
    /// Operation timeout
    Timeout,
    /// Invalid power domain
    InvalidDomain,
    /// BIU idle operation failed
    BiuIdleFailed,
    /// Power status check failed
    PowerStatusFailed,
}

impl RK3588Pmu {
    /// 创建新的 PMU 控制器实例
    ///
    /// # 参数
    /// - `base_addr`: PMU 寄存器基地址 (默认: 0xFD8D8000)
    ///
    /// # 示例
    /// ```
    /// let pmu = RK3588Pmu::new(PMU_BASE_ADDR);
    /// ```
    pub const fn new(base_addr: usize) -> Self {
        Self { base_addr }
    }

    /// 读取寄存器
    #[inline]
    fn read_reg(&self, offset: usize) -> u32 {
        unsafe { read_volatile((self.base_addr + offset) as *const u32) }
    }

    /// 写入寄存器
    #[inline]
    fn write_reg(&self, offset: usize, value: u32) {
        unsafe {
            write_volatile((self.base_addr + offset) as *mut u32, value);
        }
        compiler_fence(Ordering::SeqCst);
    }

    /// 使用写掩码机制写入寄存器
    ///
    /// RK3588 寄存器格式:
    /// - Bit [31:16]: 写掩码 (1=允许写入对应的数据位)
    /// - Bit [15:0]: 数据位
    #[allow(dead_code)]
    #[inline]
    fn write_with_mask(&self, offset: usize, mask: u32, value: u32) {
        let write_mask = mask << 16;
        self.write_reg(offset, write_mask | value);
    }

    /// 检查电源域是否开启
    ///
    /// # 参数
    /// - `domain`: 电源域
    ///
    /// # 返回
    /// - `true`: 电源开启
    /// - `false`: 电源关闭
    pub fn is_power_on(&self, domain: PowerDomain) -> bool {
        // 某些域 (如 GMAC) 使用 repair_status 而不是普通状态寄存器
        if let Some((repair_offset, repair_mask)) = domain.repair_status_register_info() {
            let status = self.read_reg(repair_offset);
            // repair_status: 1=开启, 0=关闭
            return (status & repair_mask) != 0;
        }

        // 普通状态寄存器检查
        if let Some((status_offset, mask)) = domain.status_register_info() {
            let status = self.read_reg(status_offset);
            // 状态位: 0=开启, 1=关闭
            return (status & mask) == 0;
        }

        // 如果既没有状态寄存器也没有 repair_status，检查 idle 状态
        // 对于 idle-only 域: idle=false 表示开启
        if domain.idle_register_info().is_some() {
            return !self.is_biu_idle(domain);
        }

        // 没有任何状态检查方式，假设开启
        log::warn!("Domain {:?} has no status register, assuming ON", domain);
        true
    }

    /// 检查 BIU 是否处于空闲状态
    fn is_biu_idle(&self, domain: PowerDomain) -> bool {
        if let Some((status_offset, mask)) = domain.idle_status_register_info() {
            let status = self.read_reg(status_offset);
            (status & mask) == mask
        } else {
            true // 没有 BIU 接口，认为空闲
        }
    }

    /// 检查 BIU 空闲确认状态
    fn is_biu_idle_ack(&self, domain: PowerDomain) -> bool {
        if let Some((ack_offset, mask)) = domain.idle_ack_register_info() {
            let status = self.read_reg(ack_offset);
            (status & mask) == mask
        } else {
            true // 没有 BIU 接口，认为已确认
        }
    }

    /// 发送 BIU 空闲请求
    ///
    /// # 参数
    /// - `domain`: 电源域
    /// - `idle`: true=请求空闲, false=恢复运行
    fn set_biu_idle_request(&self, domain: PowerDomain, idle: bool) -> Result<(), PmuError> {
        if let Some((idle_offset, mask, write_mask)) = domain.idle_register_info() {
            // 设置空闲请求位
            let value = if idle { mask } else { 0 };
            self.write_reg(idle_offset, write_mask | value);

            // 等待 BIU 确认和空闲状态
            for _ in 0..TIMEOUT_LOOPS {
                if self.is_biu_idle_ack(domain) && self.is_biu_idle(domain) {
                    return Ok(());
                }
                // 简单延迟
                for _ in 0..100 {
                    compiler_fence(Ordering::SeqCst);
                }
            }

            log::error!(
                "BIU idle request timeout for domain {:?}, idle={}",
                domain,
                idle
            );
            return Err(PmuError::BiuIdleFailed);
        }

        Ok(()) // 没有 BIU 接口，直接返回成功
    }

    /// 软件电源关闭流程
    ///
    /// 按照 RK3588 TRM Chapter 7.5.3 的流程:
    /// 1. 发送 BIU 空闲请求 (如果存在 BIU 接口)
    /// 2. 等待 BIU 空闲确认
    /// 3. 设置电源门控寄存器关闭电源
    /// 4. 等待电源状态确认
    ///
    /// # 参数
    /// - `domain`: 要关闭的电源域
    ///
    /// # 返回
    /// - `Ok(())`: 成功
    /// - `Err(PmuError)`: 失败
    pub fn power_off(&self, domain: PowerDomain) -> Result<(), PmuError> {
        log::info!("Power off domain: {:?}", domain);

        // 如果已经关闭，直接返回
        if !self.is_power_on(domain) {
            log::debug!("Domain {:?} is already powered off", domain);
            return Ok(());
        }

        // 步骤 1-2: 发送 BIU 空闲请求并等待确认
        self.set_biu_idle_request(domain, true)?;

        // 步骤 3: 设置电源门控寄存器关闭电源
        let (pwr_offset, mask, write_mask) = domain.register_info();
        self.write_reg(pwr_offset, write_mask | mask); // 设置位为 1 表示关闭

        // 步骤 4: 等待电源关闭状态确认
        for _ in 0..TIMEOUT_LOOPS {
            if !self.is_power_on(domain) {
                log::info!("Domain {:?} powered off successfully", domain);
                return Ok(());
            }
            // 延迟
            for _ in 0..100 {
                compiler_fence(Ordering::SeqCst);
            }
        }

        log::error!("Power off timeout for domain {:?}", domain);
        Err(PmuError::Timeout)
    }

    /// 软件电源开启流程 (无修复)
    ///
    /// 按照 RK3588 TRM Chapter 7.5.3 的流程:
    /// 1. 设置电源门控寄存器开启电源
    /// 2. 等待电源状态确认
    /// 3. 发送 BIU 恢复请求 (如果存在 BIU 接口)
    /// 4. 等待 BIU 退出空闲状态
    ///
    /// # 参数
    /// - `domain`: 要开启的电源域
    ///
    /// # 返回
    /// - `Ok(())`: 成功
    /// - `Err(PmuError)`: 失败
    pub fn power_on(&self, domain: PowerDomain) -> Result<(), PmuError> {
        log::info!("Power on domain: {:?}", domain);

        // 如果已经开启，直接返回
        if self.is_power_on(domain) {
            log::debug!("Domain {:?} is already powered on", domain);
            return Ok(());
        }

        // 步骤 1: 设置电源门控寄存器开启电源
        let (pwr_offset, _mask, write_mask) = domain.register_info();
        self.write_reg(pwr_offset, write_mask); // 清除位 (写 0) 表示开启

        // 步骤 2: 等待电源开启状态确认
        for _ in 0..TIMEOUT_LOOPS {
            if self.is_power_on(domain) {
                log::debug!("Domain {:?} power on confirmed", domain);
                break;
            }
            // 延迟
            for _ in 0..100 {
                compiler_fence(Ordering::SeqCst);
            }
        }

        if !self.is_power_on(domain) {
            log::error!("Power on timeout for domain {:?}", domain);
            return Err(PmuError::Timeout);
        }

        // 步骤 3-4: 发送 BIU 恢复请求并等待退出空闲
        self.set_biu_idle_request(domain, false)?;

        log::info!("Domain {:?} powered on successfully", domain);
        Ok(())
    }

    /// 获取电源状态
    ///
    /// # 参数
    /// - `domain`: 电源域
    ///
    /// # 返回
    /// 电源状态 (On/Off)
    pub fn get_power_state(&self, domain: PowerDomain) -> PowerState {
        if self.is_power_on(domain) {
            PowerState::On
        } else {
            PowerState::Off
        }
    }

    /// 读取 PMU 版本寄存器
    pub fn version(&self) -> u32 {
        self.read_reg(regs::offset::PMU_VERSION)
    }

    /// 读取全局电源状态寄存器
    pub fn global_power_status(&self) -> u32 {
        self.read_reg(regs::offset::PMU_GLB_POWER_STS)
    }

    /// 批量电源管理：关闭多个电源域
    ///
    /// # 参数
    /// - `domains`: 要关闭的电源域列表
    ///
    /// # 返回
    /// - `Ok(())`: 全部成功
    /// - `Err(PmuError)`: 至少有一个失败
    pub fn power_off_multiple(&self, domains: &[PowerDomain]) -> Result<(), PmuError> {
        for domain in domains {
            self.power_off(*domain)?;
        }
        Ok(())
    }

    /// 批量电源管理：开启多个电源域
    ///
    /// # 参数
    /// - `domains`: 要开启的电源域列表
    ///
    /// # 返回
    /// - `Ok(())`: 全部成功
    /// - `Err(PmuError)`: 至少有一个失败
    pub fn power_on_multiple(&self, domains: &[PowerDomain]) -> Result<(), PmuError> {
        for domain in domains {
            self.power_on(*domain)?;
        }
        Ok(())
    }

    /// 通过 CRU 全局软件复位重启系统
    ///
    /// 这是最直接、最快速的重启方式，通过向 CRU
    /// 模块的特定寄存器写入魔术值来断言系统复位信号。
    ///
    /// # 参数
    /// - `cru_base`: CRU 基地址 (默认: 0xFD7C0000)
    /// - `level`: 重启级别
    ///   - `RebootLevel::First`: 彻底复位，只保留仅支持硬件复位的寄存器
    ///   - `RebootLevel::Second`: 保留状态复位，保留 GRFs 和 GPIOs 的状态
    /// - `reset_counter_threshold`: 全局复位保持时间（以 OSC
    ///   时钟周期为单位，最长 1ms）
    ///
    /// # Safety
    /// 此函数会立即触发系统重启，调用后系统将开始重启流程。
    /// 这是一个硬件自解除的复位机制。
    ///
    /// # 示例
    /// ```no_run
    /// use pmu_rk3588::{CRU_BASE_ADDR, RK3588Pmu, RebootLevel};
    ///
    /// let pmu = RK3588Pmu::new(0xFD8D8000);
    /// // 使用第一级复位（彻底复位）
    /// unsafe {
    ///     pmu.reboot_via_cru(CRU_BASE_ADDR, RebootLevel::First, None);
    /// }
    /// ```
    pub unsafe fn reboot_via_cru(
        &self,
        cru_base: usize,
        level: RebootLevel,
        reset_counter_threshold: Option<u32>,
    ) {
        log::info!("Initiating system reboot via CRU {:?} reset", level);

        // 1. (可选) 设置全局复位保持时间
        if let Some(threshold) = reset_counter_threshold {
            let threshold = threshold & 0x3FF; // 限制在 10 位
            log::debug!(
                "Setting global reset counter threshold to {} OSC cycles",
                threshold
            );
            unsafe {
                write_volatile(
                    (cru_base + regs::cru_offset::CRU_GLB_CNT_TH) as *mut u32,
                    threshold,
                );
            }
            compiler_fence(Ordering::SeqCst);
        }

        // 2. 根据重启级别选择对应的寄存器和魔术值
        let (reset_reg, magic_value) = match level {
            RebootLevel::First => (
                cru_base + regs::cru_offset::CRU_GLB_SRST_FST_VALUE,
                regs::cru_magic::GLB_SRST_FST_MAGIC,
            ),
            RebootLevel::Second => (
                cru_base + regs::cru_offset::CRU_GLB_SRST_SND_VALUE,
                regs::cru_magic::GLB_SRST_SND_MAGIC,
            ),
        };

        log::warn!(
            "Triggering system reset by writing 0x{:04X} to 0x{:08X}",
            magic_value,
            reset_reg
        );

        // 3. 触发复位：将魔术值写入对应的复位寄存器
        // 一旦写入，硬件将断言复位信号，系统将立即开始重启
        unsafe {
            write_volatile(reset_reg as *mut u32, magic_value);
        }
        compiler_fence(Ordering::SeqCst);

        // 注意：正常情况下执行不会到达这里，因为系统已经开始重启
        log::error!("System reboot failed or delayed");

        // 无限循环等待重启生效
        loop {
            compiler_fence(Ordering::SeqCst);
        }
    }
}

/// 全局 PMU 实例 (需要在初始化时设置)
static mut GLOBAL_PMU: Option<RK3588Pmu> = None;

/// 初始化全局 PMU 实例
///
/// # Safety
/// 此函数应该在系统初始化时调用一次
pub unsafe fn init_global_pmu(base_addr: usize) {
    GLOBAL_PMU = Some(RK3588Pmu::new(base_addr));
    log::info!("RK3588 PMU initialized at 0x{:08X}", base_addr);
}

/// 获取全局 PMU 实例
///
/// # Safety
/// 必须先调用 `init_global_pmu` 初始化
pub unsafe fn get_global_pmu() -> Option<&'static RK3588Pmu> {
    unsafe { GLOBAL_PMU.as_ref() }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_power_domain_register_info() {
        let domain = PowerDomain::GPU;
        let (offset, mask, write_mask) = domain.register_info();
        assert_eq!(offset, regs::offset::PMU_PWR_GATE_SFTCON0);
        assert_eq!(mask, 1 << 0);
        assert_eq!(write_mask, 1 << 16);

        let domain = PowerDomain::VOP;
        let (offset, mask, write_mask) = domain.register_info();
        assert_eq!(offset, regs::offset::PMU_PWR_GATE_SFTCON1);
        assert_eq!(mask, 1 << 1);
        assert_eq!(write_mask, 1 << 17);
    }

    #[test]
    fn test_idle_register_info() {
        let domain = PowerDomain::GPU;
        let info = domain.idle_register_info();
        assert!(info.is_some());

        let domain = PowerDomain::GMAC;
        let info = domain.idle_register_info();
        assert!(info.is_none()); // GMAC 没有 BIU 接口
    }
}
